package com.forum.config.WebSocket;


import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.forum.entity.RequestData;
import com.forum.service.impl.AIQuestionService;
import com.fasterxml.jackson.databind.ObjectMapper;

import io.reactivex.Flowable;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;


@Service
@Slf4j
public class MyWebSocketHandler extends TextWebSocketHandler {
    private final AIQuestionService aiQuestionService;
    private final ObjectMapper objectMapper = new ObjectMapper();

    @Autowired
    public MyWebSocketHandler(AIQuestionService aiQuestionService) {
        this.aiQuestionService = aiQuestionService;
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        RequestData requestData = objectMapper.readValue(message.getPayload(), RequestData.class);
        String question = requestData.getQuestion();
        int userId = requestData.getUserId();

        // 调用askAIWithContext方法并获取回答的流式响应
        Flowable<GenerationResult> result = aiQuestionService.askAIWithContext(question, userId);

        result.blockingForEach(genResult -> {
            String content = genResult.getOutput().getChoices().get(0).getMessage().getContent();
            try {
                session.sendMessage(new TextMessage(content));
            } catch (Exception e) {
                log.error("通过 WebSocket 发送消息时出错: {}", e.getMessage());
            }
        });
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        System.out.println("已建立 WebSocket 连接: " + session.getId());
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        System.out.println("WebSocket 连接已关闭: " + session.getId());
    }
}